import os
import torch
import math
import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd

class Experiment:
    """
    implement the experiment and responding result saving and showing.
    """
    def __init__(self, name, config):

        self.root_path = './Result'
        self.name = name
        self.config = config
        self.basic_path = os.path.join(self.root_path, self.name)
        if os.path.exists(self.basic_path) == False:
            os.makedirs(self.basic_path)
    
    def generate_name(self, case='Case1'):

        father_path = '{}/{}/{}/Batch({})_Size({})_Mode({})'.format(case, self.config['model_arch'], self.config['dataset'], self.config['base_batch'], self.config['base_size'], self.config['base_mode'])
        path = os.path.join(self.basic_path, father_path, 'Eps({:.4f})_Norm({})_Method({})'.format(self.config['epsilon'], self.config['attck_norm'], self.config['attck_method']))
        
        return path
    
    def load_file(self, case='Case1'):

        if os.path.exists(os.path.join(self.generate_name(case=case), 'result.pt')) == False:
            raise Exception('Do not find result in path {}'.format(os.path.join(self.generate_name(case=case), 'result.pt')))

        dicter = torch.load(os.path.join(self.generate_name(case=case), 'result.pt'))

        return dicter

    def save_total_file(self, case='Case1'):

        if os.path.exists(os.path.join(self.basic_path, '{}'.format(case))) == False:
            raise Exception('Did not have the following path <{}>'.format(os.path.join(self.basic_path, '{}'.format(case))))
        
        if os.path.exists(os.path.join(self.basic_path, '{}'.format(case), 'data.pt')) == False:
            dicter = {}
        else:
            dicter = torch.load(os.path.join(self.basic_path, '{}'.format(case), 'data.pt'))

        dataset_name = self.config['dataset']
        model_arch = self.config['model_arch']

        if dataset_name not in dicter.keys():
            dicter[dataset_name] = {}
        
        sub_dicter = dicter[dataset_name]

        context = self.load_file(case=case)

        sub_dicter[model_arch] = context
    
        torch.save(dicter, os.path.join(self.basic_path, '{}'.format(case), 'data.pt'))
    
    def load_total_file(self, case='Case1'):

        if os.path.exists(os.path.join(self.basic_path, '{}'.format(case))) == False:
            raise Exception('Did not have the following path <{}>'.format(os.path.join(self.basic_path, '{}'.format(case))))
        
        if os.path.exists(os.path.join(self.basic_path, '{}'.format(case), 'data.pt')) == False:
            raise Exception('No such directory or file, please recheck.')
        else:
            dicter = torch.load(os.path.join(self.basic_path, '{}'.format(case), 'data.pt'))

        return dicter
    
    def latex_export(self, case='Case1'):
        
        ## for robust model case
        if self.config['dataset'] == 'adv_ResNet18':
            
            raise Exception('The response code has not been perfected yet')
            pass

        ## for target and untarget experiment
        else:

            row_list = [
                ['MNIST', 'CNN'], 
                ['FMNIST', 'CNN'], 
                ['CIFAR-10', 'VGG16'], 
                ['CIFAR-10', 'ResNet18'],
            ]

            if case == 'Case1':
                
                attack_method_list = ['OnePixel', 'SimBA', 'FNS', 'MPack', 'GA', 'SMPack', 'FGSM', 'PGD']
                dicter = self.load_total_file(case='Case1')

                print('CASE: 1')
                ## print(dicter)

                for row in row_list:
                    print('Dataset: {}, Model: {}'.format(row[0], row[1]))
                    str = ''
                    sub_dicter = dicter[row[0]][row[1]]
                    for attack_method in attack_method_list:
                        
                        if attack_method in sub_dicter.keys():
                            str += '&${:.1f}\\pm{:.1f}$ '.format(sub_dicter[attack_method][0] * 100, math.sqrt(sub_dicter[attack_method][1]) * 100)
                        else:
                            str += '&- '
                    str += '\\\\'
                    print(str)
                    print()
        
            elif case == 'Case2':

                attack_method_list = ['OnePixel', 'SimBA', 'FNS', 'MPack', 'GA', 'SMPack']
                reorganize_dicter = {
                    'average': [[] for i in attack_method_list],
                    'median': [[] for i in attack_method_list],
                }
                dicter = self.load_total_file(case='Case2')

                print('CASE: 2')
                ## print(dicter)

                for row in row_list:
                    print('Dataset: {}, Model: {}'.format(row[0], row[1]))
                    avg_str = ''
                    median_str = ''
                    sub_dicter = dicter[row[0]][row[1]]
                    for dex, attack_method in enumerate(attack_method_list):
                        
                        if attack_method in sub_dicter.keys():
                            avg_str += '&${:.1f}\\pm{:.1f}$ '.format(sub_dicter[attack_method]['Average'][0], math.sqrt(sub_dicter[attack_method]['Average'][1]))
                            reorganize_dicter['average'][dex].append('&${:.1f}\\pm{:.1f}$ '.format(sub_dicter[attack_method]['Average'][0], math.sqrt(sub_dicter[attack_method]['Average'][1])))
                            median_str += '&${:.1f}\\pm{:.1f}$ '.format(sub_dicter[attack_method]['Median'][0], math.sqrt(sub_dicter[attack_method]['Median'][1]))
                            reorganize_dicter['median'][dex].append('&${:.1f}\\pm{:.1f}$ '.format(sub_dicter[attack_method]['Median'][0], math.sqrt(sub_dicter[attack_method]['Median'][1])))
                        else:
                            avg_str += '&- '
                            reorganize_dicter['average'][dex].append('&- ')
                            median_str += '&- '         
                            reorganize_dicter['median'][dex].append('&- ')
                        
                    avg_str += '\\\\'                                        
                    median_str += '\\\\'
        
                    print(avg_str)
                    print()
                    print(median_str)

                    print('\n')

                for dex in range(len(reorganize_dicter['average'])):
                    reorganize_dicter['average'][dex].append('\\\\')
                    reorganize_dicter['median'][dex].append('\\\\')

                print('Type 2 show:')

                print('Average:')
                for lister in reorganize_dicter['average']:
                    
                    for dex, sub_str in enumerate(lister):
                        if dex == 1 or dex == 2:
                            print('& ', end='')
                        print(sub_str, end='')
                    print()
                
                print('Median')
                for lister in reorganize_dicter['median']:
                    
                    for dex, sub_str in enumerate(lister):
                        if dex == 1 or dex == 2:
                            print('& ', end='')
                        print(sub_str, end='')
                    print()

            else:

                raise ValueError

    def show_result(self, case='Case1'):

        print(self.load_file(case=case)) 
    
    
    def base_size_plot(self, max_size=10):

        sns.set_style('whitegrid')
        fig = plt.figure(figsize=(10, 8))


        for index, model_arch in enumerate(['VGG16', 'ResNet18']):
            
            self.config['model_arch'] = model_arch
            ax = plt.subplot(2, 1, index + 1)

            base_size1 = []
            value1 = []

            base_size2 = []
            value2 = []

            for size in range(1, max_size + 1):
                
                self.config['base_size'] = size
                base_size1.append(size)
                temp_dicter = self.load_file(case='Case1')
                value1.append(temp_dicter['MPack'][0])
            
            for size in range(1, max_size + 1):
                
                self.config['base_size'] = size
                base_size2.append(size)
                temp_dicter = self.load_file(case='Case2')
                value2.append(temp_dicter['SMPack']['Average'][0])

            ax.plot(base_size1, value1, linewidth=2.3, marker='v', label='MPack')
            ax.legend(loc=2, fontsize=16)

            ## plt.ylim(min(value) - 0.1, max(value) + 0.03)
            plt.ylim(min(value1) - 0.07, max(value1) + 0.05)
            plt.xlabel('MUAPs  size', fontsize=18)
            plt.ylabel('ASR', fontsize=18)

            plt.xticks(fontsize=14)
            plt.yticks(fontsize=14)

            ax2 = ax.twinx()
            ax2.plot(base_size2, value2, linewidth=2.3, marker='^', color='#C44E52', label='SMPack')
            ax2.set_ylabel('AvgQue', fontsize=18)
            ax2.legend(loc=1, fontsize=16)
            plt.ylim(min(value2) - 50, max(value2) + 100)

            plt.title('{}  '.format(self.config['model_arch']), fontsize=20)

        plt.subplots_adjust(wspace=0.2, hspace=0.4)
        plt.savefig('./img/param_Base_size.pdf'.format(self.config['model_arch']), dpi=800, bbox_inches='tight', pad_inches=0.2)
        print('Save done')
        plt.close()    


    def base_batch_plot(self, ):

        sns.set_style('whitegrid')
        fig = plt.figure(figsize=(10, 8))


        for index, model_arch in enumerate(['VGG16', 'ResNet18']):
            
            self.config['model_arch'] = model_arch
            ax = plt.subplot(2, 1, index + 1)

            base_batch1 = []
            value1 = []

            base_batch2 = []
            value2 = []

            for batch in [1, 2, 4, 8, 16]:
                
                self.config['base_batch'] = batch
                base_batch1.append(batch)
                temp_dicter = self.load_file(case='Case1')
                value1.append(temp_dicter['MPack'][0])
            
            for batch in [1, 2, 4, 8, 16]:
                
                self.config['base_batch'] = batch
                base_batch2.append(batch)
                temp_dicter = self.load_file(case='Case2')
                value2.append(temp_dicter['SMPack']['Average'][0])


            ax.plot(base_batch1, value1, linewidth=2.3, marker='v', label='MPack')
            ax.legend(loc=2, fontsize=16)

            ## plt.ylim(min(value) - 0.1, max(value) + 0.03)
            plt.ylim(min(value1) - 0.07, max(value1) + 0.05)
            plt.xlabel('MUAPs  batch', fontsize=18)
            plt.ylabel('ASR', fontsize=18)

            plt.xticks(fontsize=14)
            plt.yticks(fontsize=14)

            ax2 = ax.twinx()
            ax2.plot(base_batch2, value2, linewidth=2.3, marker='^', color='#C44E52', label='SMPack')
            ax2.set_ylabel('AvgQue', fontsize=18)
            ax2.legend(loc=1, fontsize=16)
            plt.ylim(min(value2) - 50, max(value2) + 100)

            plt.title('{}  '.format(self.config['model_arch']), fontsize=20)

        plt.subplots_adjust(wspace=0.2, hspace=0.4)
        plt.savefig('./img/param_Base_batch.pdf'.format(self.config['model_arch']), dpi=800, bbox_inches='tight', pad_inches=0.2)
        print('Save done')
        plt.close()      